{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Pulse Optimization Using GRAPE \n",
    "\n",
    "*Copyright (c) 2021 Institute for Quantum Computing, Baidu Inc. All Rights Reserved.*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Outline\n",
    "This tutorial introduces how to implement the single-qubit gate using the GRAPE (GRadient Ascent Pulse Engineering) algorithm in Quanlse. The outline of this tutorial is as follows:\n",
    "- Introduction\n",
    "- Preparation\n",
    "- Construct Hamiltonian\n",
    "- Generate optimized pulses using GRAPE\n",
    "- Summary"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Introduction\n",
    "\n",
    "**Mechanism**\n",
    "\n",
    "Gradient-based algorithms are expected to search for the extreme value of an objective function. To search for the maximum value, we can use the gradient ascent algorithm: \n",
    "$$\n",
    "\\Theta_n = \\Theta_{n-1} + k \\bigtriangledown J(\\Theta), \\tag{1}\n",
    "$$\n",
    "\n",
    "The gradient ascent algorithm is analogous to the process of climbing a mountain. If we want to ascent to the peak, we need to know each step's length and direction which give $k\\bigtriangledown{J(\\Theta)}$ when multiplied. According to the equation above, the position after each step $\\Theta_n$ is determined by the product of the derivative of the original position $\\Theta_{n-1}$, the objective function $J(\\Theta)$ at the original position, and the coefficient $k$. \n",
    "\n",
    "\n",
    "The GRAPE algorithm \\[1\\] is a prototypical method in quantum optimal control - initially introduced on the NMR platform and later extended to other platforms. Experimentally, quantum gates are implemented by external control fields in form of pulse or magnetic flux. By chopping the total control time (for quantum gates) into intervals and assuming the amplitudes of each slice are constant, the amplitudes of the pulses in each slice can be taken as the parameters for optimization. In the Heisenberg picture, we get operator $U (t)$ for dynamics in the period $t \\in \\left[0, T\\right]$：\n",
    "$$\n",
    "i\\hbar \\frac{\\partial U (t)}{\\partial t} = \\hat{H}(t)U(t),\n",
    "\\tag{2}\n",
    "$$\n",
    "\n",
    "where $\\hbar$ is the reduced Planck constant (here we choose our units so that $\\hbar=1$). When the Hamiltonian is time-independent，the solution $U (t = T)$ is\n",
    "$$\n",
    "U(T) = \\exp(-i\\hat{H}T).\n",
    "\\tag{3}\n",
    "$$\n",
    "\n",
    "For a time-dependent Hamiltonian, an effective method is to chop the total time $T$ into intervals of the same length, assuming that the Hamiltonian $H_{j}$ is time-independent for each time period $\\left[t_j, t_{j+1}\\right]$. Here the subscript $j$ denotes the slice sequence. We can solve the respective dynamics $U_{j}$ for each time slice using Eq. (2). The whole process can be represented as:\n",
    "$$\n",
    "U = U_{N}U_{N-1}\\cdots{U}_{1},\n",
    "\\tag{4}\n",
    "$$\n",
    "\n",
    "where $N$ denotes the number of slices."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Implementation**\n",
    "\n",
    "Here, we introduce a GRAPE process and set the square of the unnormalized gate fidelity as the objective function $J$:\n",
    "$$\n",
    "J = \\left|{\\rm Tr}\\left[(U^{\\dagger}_{\\rm target}U(T)\\right]\\right|^2 = \\left|{\\rm Tr}\\left[P^{\\dagger}_{j}X_{j}\\right]\\right|^2,  \\ \\forall\\ 0<j<N, \\tag{5}\n",
    "$$\n",
    "\n",
    "where $U_{\\rm target}$ is the target unitary matrix; $X_{j} = U_{j}\\cdots{U}_{1}$ is the intermediate propagator; and $P_{j} = U^{\\dagger}_{j+1}\\cdots{U}^{\\dagger}_{N}U_{\\rm target}$ is the intermediate back-propagator.  \n",
    "\n",
    "Continous pulses $u_i(t)$ on the control Hamiltonian $\\hat{H}_{i}$ are chopped into discrete functions $u_i(j)$. The partial derivative of the objective function $J$ is given \\[2\\]:\n",
    "$$\n",
    "\\dfrac{\\partial J}{\\partial u_i(j)} = {\\rm Re}\\{-2 i \\Delta t \\left[{\\rm Tr}(P^{\\dagger}_{j}\\hat{H}_iX_{j}){\\rm Tr}(P^{\\dagger}_{j}X_{j})\\right]\\},\n",
    "\\tag{6}\n",
    "$$\n",
    "\n",
    "Gradient-based algorithms will reshape the pulse after each iteration with the learning rate $k$ fixed. Finally, we can obtain pulses with optimized shapes:\n",
    "$$\n",
    "u_i(t) \\mapsto u_i(t) + k\\dfrac{\\partial J}{\\partial u_i(j)},\\tag{7}\n",
    "$$\n",
    "\n",
    "\n",
    "![GRAPE](figures/GRAPE1.png)\n",
    "\n",
    "As shown in the figure above, GRAPE reshapes the pulse's waveform after each iteration.\n",
    "\n",
    "In the following text, we demonstrate how one can generate optimized pulse sequence for single-qubit gate using GRAPE."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Preparation\n",
    "\n",
    "After you have successfully installed Quanlse on your device, you could run the Quanlse program below following this tutorial. To run this particular tutorial, you would need to import the following packages from Quanlse and other commonly-used Python libraries:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import required packages\n",
    "from Quanlse.remoteOptimizer import remoteOptimize1QubitGRAPE\n",
    "from Quanlse.QOperator import number, driveX, driveY, duff\n",
    "from Quanlse.QHamiltonian import QHamiltonian as QHam\n",
    "from Quanlse.QOperation import FixedGate\n",
    "from Quanlse.Utils.Functions import project\n",
    "from Quanlse.QWaveform import gaussian\n",
    "\n",
    "import numpy as np\n",
    "from numpy import dot, round"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, we define constants that describe the overall system, including the sampling period of the arbitrary wave generator (AWG), the system's energy levels, and the gate duration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Sampling period\n",
    "dt = 1.0\n",
    "\n",
    "# System energy level\n",
    "level = 2\n",
    "\n",
    "# Duration of the gate (ns)\n",
    "tg = 40"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Construct Hamiltonian\n",
    "\n",
    "In this section, we will construct the Hamiltonian. For a two-level system, the Hamiltonian in the rotating frame is made up by：\n",
    "$$\n",
    "\\hat{H} = \\frac{1}{2} \\Omega^x(t) (\\hat{a}+\\hat{a}^{\\dagger}) + i \\frac{1}{2} \\Omega^y(t) (\\hat{a}-\\hat{a}^{\\dagger}),\n",
    "$$\n",
    "\n",
    "$\\Omega^x(t)$ and $\\Omega^y(t)$ are the control pulses on the X channel and Y channel respectively. We construct the Hamiltonian as follows："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize the Hamiltonian\n",
    "ham = QHam(subSysNum=1, sysLevel=level, dt=dt)\n",
    "\n",
    "# Add the anharmonicity term\n",
    "alphaQ = - 0.22 * (2 * np.pi)\n",
    "ham.addDrift(duff, 0, coef=alphaQ)\n",
    "\n",
    "# Add the control terms\n",
    "ham.addWave(driveX, 0, waves=gaussian(tg, a=0.3, tau=tg / 2, sigma=tg / 8))\n",
    "ham.addWave(driveY, 0, waves=gaussian(tg, a=0.3, tau=tg / 2, sigma=tg / 8))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The optimization usually takes a long time to process on local devices, however, we provide a cloud service that could speed up this process significantly. To use Quanlse Cloud Service, the users need to acquire a token from http://quantum-hub.baidu.com and use the following command to submit the job onto Quanlse's server. For this example, we can write:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Import Define class and set the token\n",
    "# Please visit http://quantum-hub.baidu.com\n",
    "from Quanlse import Define\n",
    "Define.hubToken = ''"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generate Optimized Pulses Using GRAPE\n",
    "\n",
    "Using the function `remoteOptimize1QubitGRAPE()`, we can obtain the optimized pulse on each channel and the infidelity of the unitary evolution. The function `remoteOptimize1QubitGRAPE()` takes in a Hamiltonian dictionary `ham`, a target unitary matrix `ugoal`, the time duration of the gate (`tg`, default at 20 nanoseconds), the maximum iteration (`iterate`, default at 150), and a list consisting the pulse numbers for each channel (`xyzPulses`, default at \\[1, 1, 0\\]). The number of slices is determined by $tg / dt$. Here, we set our goal unitary evolution to be the H gate:\n",
    "$$\n",
    "U_{\\rm target} = \n",
    "\\dfrac{1}{\\sqrt{2}}\\begin{bmatrix} \n",
    "1 & 1  \\\\ 1 & -1 \n",
    "\\end{bmatrix} .\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the target unitary evolution\n",
    "uGoal = FixedGate.H.getMatrix()\n",
    "\n",
    "# Run the optimization\n",
    "job, infid = remoteOptimize1QubitGRAPE(ham, uGoal, tg=tg, iterate=50, xyzPulses=None)\n",
    "\n",
    "# Print infidelity and the waveforms\n",
    "print(f\"minimum infidelity with GRAPE: {infid}\")\n",
    "ham.plot(color = ['blue', 'mint'], dark=True)\n",
    "\n",
    "# Print the evolution\n",
    "result = ham.simulate()\n",
    "print(\"The evolution U:\\n\", round(result.result[0][\"unitary\"], 2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The quantum gate optimized by the GRAPE algorithm usually differs from the target gate by a global phase $e^{i\\phi}$, which will not affect the result."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "This tutorial introduces how to generate optimized single-qubit gate pulses using the GRAPE algorithm. We can see that GRAPE generate pulses with somewhat irregular shapes. As one of the most common optimization algorithms, GRAPE converges quickly when given a reasonable initial value. The users are encouraged to try parameter values different from this tutorial to obtain the optimal result. (try varying the sampling size when result is undesirable)\n",
    "\n",
    "The users could follow this link [tutorial-GRAPE.ipynb](https://github.com/baidu/Quanlse/blob/main/Tutorial/EN/tutorial-GRAPE.ipynb) to the GitHub page of this Jupyter Notebook document to download the code above and further explore the GRAPE algorithm for other gates."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## References\n",
    "\n",
    "\\[1\\] [Wilhelm, Frank K., et al. \"An introduction into optimal control for quantum technologies.\" *arXiv preprint arXiv*:2003.10132 (2020).](https://arxiv.org/abs/2003.10132v1)\n",
    "\n",
    "\\[2\\] [Khaneja, Navin, et al. \"Optimal control of coupled spin dynamics: design of NMR pulse sequences by gradient ascent algorithms.\" *Journal of magnetic resonance* 172.2 (2005): 296-305.](https://doi.org/10.1016/j.jmr.2004.11.004)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
